#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RSAAlgorithm的泄漏模型
定义RSA各个操作的Power泄漏特征
"""

import numpy as np
from typing import Dict, Any, List


class RSALeakageModels:
    """RSA泄漏模型集合"""
    
    @staticmethod
    def hamming_weight(value: int) -> int:
        """Calculate汉明权重"""
        return bin(value).count('1')
    
    @staticmethod
    def hamming_distance(a: int, b: int) -> int:
        """CalculateHamming distance"""
        return RSALeakageModels.hamming_weight(a ^ b)
    
    @classmethod
    def modular_multiplication_model(cls, a: int, b: int, n: int) -> int:
        """
        模乘运算泄漏模型
        
        Args:
            a: 乘数a
            b: 乘数b
            n: 模数
            
        Returns:
            模乘Result (a * b) mod n
        """
        return (a * b) % n
    
    @classmethod
    def modular_exponentiation_model(cls, base: int, exponent: int, modulus: int) -> int:
        """
        模幂运算泄漏模型（简化版）
        
        Args:
            base: 底数
            exponent: 指数
            modulus: 模数
            
        Returns:
            模幂Result base^exponent mod modulus
        """
        # Use快速模幂Algorithm
        result = 1
        base = base % modulus
        
        while exponent > 0:
            if exponent & 1:
                result = (result * base) % modulus
            exponent >>= 1
            base = (base * base) % modulus
        
        return result
    
    @classmethod
    def square_and_multiply_model(cls, base: int, bit_position: int, 
                                current_result: int, modulus: int) -> tuple:
        """
        平方乘Algorithm泄漏模型
        
        Args:
            base: 底数
            bit_position: CurrentProcess的位位置
            current_result: Current累积Result
            modulus: 模数
            
        Returns:
            (new_result, operation_type) - 新ResultAnd操作Class型
        """
        # 总YesExecute平方操作
        squared = (current_result * current_result) % modulus
        
        # IfCurrent位For1，还要Execute乘法操作
        if bit_position == 1:
            multiplied = (squared * base) % modulus
            return multiplied, "square_and_multiply"
        else:
            return squared, "square_only"
    
    @classmethod
    def montgomery_ladder_model(cls, base: int, bit: int, 
                              r1: int, r2: int, modulus: int) -> tuple:
        """
        Montgomery阶梯Algorithm泄漏模型
        
        Args:
            base: 底数
            bit: CurrentProcess的位（0Or1）
            r1: 寄存器1
            r2: 寄存器2
            modulus: 模数
            
        Returns:
            (new_r1, new_r2) - 新的寄存器值
        """
        if bit == 0:
            new_r2 = (r1 * r2) % modulus
            new_r1 = (r1 * r1) % modulus
        else:
            new_r1 = (r1 * r2) % modulus
            new_r2 = (r2 * r2) % modulus
        
        return new_r1, new_r2
    
    @classmethod
    def crt_model(cls, message: int, dp: int, dq: int, p: int, q: int) -> int:
        """
        中国剩余定理(CRT)泄漏模型
        
        Args:
            message: 消息
            dp: d mod (p-1)
            dq: d mod (q-1)
            p: 素数p
            q: 素数q
            
        Returns:
            CRTCalculateResult
        """
        # Calculate m1 = message^dp mod p
        m1 = pow(message, dp, p)
        
        # Calculate m2 = message^dq mod q
        m2 = pow(message, dq, q)
        
        # Calculate h = (m1 - m2) * qinv mod p
        qinv = pow(q, -1, p)  # q的模p逆元
        h = ((m1 - m2) * qinv) % p
        
        # CalculateFinal result m = m2 + h * q
        result = m2 + h * q
        
        return result
    
    @classmethod
    def blinding_model(cls, message: int, r: int, e: int, n: int) -> tuple:
        """
        盲化技术泄漏模型
        
        Args:
            message: 原始消息
            r: 随机盲化因子
            e: 公钥指数
            n: 模数
            
        Returns:
            (blinded_message, unblinding_factor) - 盲化消息And去盲化因子
        """
        # 盲化：m' = m * r^e mod n
        blinded_message = (message * pow(r, e, n)) % n
        
        # 去盲化因子：r^(-1) mod n
        unblinding_factor = pow(r, -1, n)
        
        return blinded_message, unblinding_factor
    
    @classmethod
    def timing_attack_model(cls, exponent_bit: int, base_hamming_weight: int) -> float:
        """
        时序Attack泄漏模型
        
        Args:
            exponent_bit: 指数位（0Or1）
            base_hamming_weight: 底数的汉明权重
            
        Returns:
            预期Execute时间（相对值）
        """
        # 基础时间（平方操作）
        base_time = 1.0
        
        # If指数位For1，Require额外的乘法操作
        if exponent_bit == 1:
            # 乘法时间与汉明权重相关
            multiply_time = 0.5 + 0.1 * base_hamming_weight
            return base_time + multiply_time
        else:
            return base_time
    
    @classmethod
    def power_analysis_model(cls, operation_type: str, operand_hw: int) -> float:
        """
        Power Analysis泄漏模型
        
        Args:
            operation_type: 操作Class型（'square', 'multiply', 'add', 'subtract'）
            operand_hw: 操作数汉明权重
            
        Returns:
            预期Power值（相对值）
        """
        # 基础Power
        base_power = {
            'square': 2.0,
            'multiply': 1.5,
            'add': 0.5,
            'subtract': 0.5
        }
        
        # Power与汉明权重成正比
        power = base_power.get(operation_type, 1.0)
        power += 0.1 * operand_hw
        
        return power
    
    @classmethod
    def get_intermediate_functions(cls) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return {
            'modular_multiplication': cls.modular_multiplication_model,
            'modular_exponentiation': cls.modular_exponentiation_model,
            'square_and_multiply': cls.square_and_multiply_model,
            'montgomery_ladder': cls.montgomery_ladder_model,
            'crt': cls.crt_model,
            'blinding': cls.blinding_model,
        }
    
    @classmethod
    def get_leakage_models(cls) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return {
            'value': lambda x: x,
            'hamming_weight': cls.hamming_weight,
            'hamming_distance': cls.hamming_distance,
            'bit_value': lambda x, bit_pos: (x >> bit_pos) & 1,
        }
    
    @classmethod
    def get_target_operations(cls) -> List[str]:
        """获取可Attack的目标操作"""
        return [
            'modular_exponentiation',
            'square_and_multiply',
            'montgomery_ladder',
            'crt_computation',
            'modular_multiplication',
            'key_generation',
            'blinding_operation',
        ]
    
    @classmethod
    def get_operation_to_intermediate_mapping(cls) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return {
            'modular_exponentiation': 'modular_exponentiation',
            'square_and_multiply': 'square_and_multiply',
            'montgomery_ladder': 'montgomery_ladder',
            'crt_computation': 'crt',
            'modular_multiplication': 'modular_multiplication',
            'key_generation': 'modular_exponentiation',
            'blinding_operation': 'blinding',
        }
    
